package tools;

import burp.IBurpCollaboratorClientContext;
import burp.IBurpCollaboratorInteraction;
import burp.IBurpExtenderCallbacks;
import org.apache.commons.lang3.StringUtils;

import java.io.UnsupportedEncodingException;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Author: 4ut15m.4ny0neSec, m4n1c.4ny0neSec
 * Date: 2022/1
 * Desc: 使用burpsuite官方的dns记录域名（等同于dnslog）
 * */
public class DNSMonitor {
    // Collaborator的客户端上下文对象
    public static IBurpCollaboratorClientContext context;
    // 在多线程使用DNSMonitor的时候，用allDNS来存放不同线程各自的dns域名
    public HashMap<String, String> allDNS;
    // 需要使用base64解码的属性
    public String[] needDecode = new String[]{"raw_query","request"};

    /**
     * Author: 4ut15m.4ny0neSec
     * Desc: 初始化
     */
    public DNSMonitor(IBurpExtenderCallbacks callbacks) {
        context = callbacks.createBurpCollaboratorClientContext();
        allDNS = new HashMap<String, String>();
    }

    /**
     * Author: 4ut15m.4ny0neSec
     * Desc: 生成一个dns记录的域名
     * */
    public String generateDNS(){
       return generateDNS("4ny0ne.one");
    }

    /**
     * Desc: 为多线程做考虑
     * Param:
     *      String threadName: 线程名，用来将多个线程获取的dns进行记录存储
     * return:
     *      String dns: 返回本次获得的dns域名
     * */
    public String generateDNS(String threadName){
        String dns = context.generatePayload(true);
        if(threadName.equals("4ny0ne.one")){
            // 不是多线程的话，就不记录dns
            return dns;
        }
        if(!allDNS.containsKey(threadName)){
            allDNS.put(threadName, dns);
        }

        return dns;
    }

    /**
     * Author: 4ut15m.4ny0neSec
     * Desc: 拉取dns收到的所有结果
     * */
    public String pullAllResults(String dns){
        return pullAllResults(dns,false)[0];
    }

    /**
     * Author: 4ut15m.4ny0neSec
     * Desc: 当retList为true时，代表将拉取的结果分开存储
     * 一个dns域名可能会记录多个请求，所以会有多个结果，如果retList为false则将
     * 所有请求的结果拼接在一起以换行符连接，而后返回
     * */
    public String[] pullAllResults(String dns, boolean retList){
        List<IBurpCollaboratorInteraction> interactions = context.fetchCollaboratorInteractionsFor(dns);
        Object[] objects = interactions.toArray();

        if(objects.length == 0){
            return new String[]{"Nothing was recorded"};
        }
        if(retList){
            List<String> results = new LinkedList<>();
            // 遍历dns请求
            for (IBurpCollaboratorInteraction interaction : interactions) {
                // 获取结果，并存放至List中
                results.add(interaction.getProperties().toString());
            }
            return (String[]) results.toArray();
        }

        String res = "";
        for (IBurpCollaboratorInteraction interaction : interactions) {
               res += interaction.getProperties().toString()+"\r\n";
        }

        return new String[]{res};
    }

    /**
     * Author: 4ut15m.4ny0neSec
     * Desc: 用来获取指定属性的值，比如client_ip
     * 一个记录包含了多个属性，样例如下：
     * {time_stamp=2022-Jan-18 17:27:19 UTC, interaction_type=collaborator, client_ip=122.228.2.4, query_type=AAAA, interaction_id=bjyefs0zpgrtuhj9moesdulq7hd71w, type=DNS, raw_query=ThgAAAABAAAAAAAAHmJqeWVmczB6cGdydHVoajltb2VzZHVscTdoZDcxdxBidXJwY29sbGFib3JhdG9yA25ldAAAHAAB}
     *
     * Params:
     *      String attribute: 指明要获取的属性，比如client_ip
     *      String dns: 指明要获取哪个dns域名的记录
     * return:
     *      String result: 经过base64解码后的指定属性值
     *
     * */
    public String takeSomeAttribute(String attribute, String dns){
        String result = "";
        boolean flag = false;
        List<IBurpCollaboratorInteraction> interactions = context.fetchCollaboratorInteractionsFor(dns);
        if(interactions.toArray().length==0){
            return "Nothing was recorded";
        }

        try {
            for (String s : needDecode) {
                if(attribute.equals(s)){
                    // 说明用户要获取的属性值是经过base64编码的，我们要对这个属性值进行解码处理
                    flag = true;
                    break;
                }
            }

            if(flag){
                for (IBurpCollaboratorInteraction interaction : interactions) {
                    String tmp = interaction.getProperty(attribute);
                    if (tmp != null){
                        result += resDecode(tmp)+"\r\n";
                    }
                }
            }else {
                for (IBurpCollaboratorInteraction interaction : interactions) {
                    String tmp = interaction.getProperty(attribute);
                    if(tmp != null){
                        result += interaction.getProperty(attribute)+"\r\n";
                    }
                }
            }

        }catch (Exception e){
            return e.getMessage();
        }

        return result;
    }

    /**
     * Author: 4ut15m.4ny0neSec
     * Desc: 用来将读取到的数据进行base64解码
     * 与parse方法搭配使用，专门用来处理collaborator返回的数据
     * Param:
     *      String encode: 经过Base64加密的数据（这里一般是collaborator返回字段的具体值）
     *
     * */
    public static String resDecode(String encode) throws UnsupportedEncodingException {
        byte[] bytes = Base64.getDecoder().decode(encode);
        String result = "";
        for (byte aByte : bytes) {
            if((aByte>=32 && aByte<=126)||aByte ==10){
                result += new String(new byte[]{aByte},"UTF-8");
            }else {
                result += ".";
            }
        }
        return parse(result);
    }

    /**
     * Author: m4n1c.4ny0neSec
     * Date:
     * Param:
     *      String input: 待解析的字符串
     * return:
     *      String res: 正则处理过后的字符串
     *
     * input传入的字符串样例：
     * .x...........4ut15m.jl818qb3x2808wvpz5p3d1p0trzhn6.burpcollaborator.net.......)........
     * JL...........6kryetcfvp048c4czps9prisrjx9ly.burpcollaborator.net.....
     *
     * 处理过后的res结果为：
     * 4ut15m.jl818qb3x2808wvpz5p3d1p0trzhn6.burpcollaborator.net
     * 6kryetcfvp048c4czps9prisrjx9ly.burpcollaborator.net
     */
    public static String parse(String input) {
        String result = "";

        String initialValue = ".burpcollaborator.net";
        //递归获取
        result = getAppend(input,initialValue);
        if(StringUtils.isNotEmpty(result)){
            return result;
        }else {
            return "Something error!";
        }

    }

    /**
     * Author: m4n1c.4ny0neSec
     * Desc: 递归获取结果字符串
     * Param:
     *      String initialValue: 递归初始确定字符串
     *      String input: 原始输入字符串
     *
     * */
    public static String getAppend(String input, String initialValue) {
        Pattern valuePattern = Pattern.compile("\\w+"+initialValue);
        Matcher valueMatcher = valuePattern.matcher(input);
        if(valueMatcher.find()){
            return getAppend(input,"." + valueMatcher.group());
        }else {
            if(initialValue.equals(".burpcollaborator.net")){
                return null;
            }else {
                return initialValue.substring(1);
            }
        }
    }


}